package com.feriki.msync.process.sink;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.feriki.msync.constants.PolicyType;
import com.feriki.msync.meta.FieldTypeMeta;
import com.feriki.msync.meta.MilvusCollectionMeta;
import com.feriki.msync.utils.IDGenerator;
import com.feriki.msync.utils.MilvusOperatorUtil;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import com.feriki.msync.constants.ConnectType;
import com.feriki.msync.pool.MilvusConnectionPool;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.grpc.MutationResult;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.CollectionSchemaParam;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.partition.CreatePartitionParam;
import io.milvus.param.partition.HasPartitionParam;
import io.milvus.response.DescCollResponseWrapper;
import lombok.extern.slf4j.Slf4j;

import java.math.BigDecimal;
import java.util.*;

/**
 * @author luozhida
 * @create 2024-07-04
 */
@Slf4j
public class MilvusSinkHandler extends SinkHandler{

    private Random random = new Random();
    private static Map<Integer,String> moniPartitionValueMap = new HashMap<Integer,String>(){
        {
            put(0,"20241231");
            put(1,"20241230");
            put(2,"20241229");
            put(3,"20241228");
            put(4,"20241227");
        }
    };

    private MilvusConnectionPool sinkPool;
    private MilvusServiceClient sinkMilvusClient;
    private final boolean enableNonAutoIdGenerateValue;

    private boolean isAutoID;
    private String primaryKey;
    private int primaryDataType;
    private String vectorKey;
    private boolean testPartition;

    public MilvusSinkHandler(MilvusConnectionPool sinkPool, String collectionName, String policy,
                             boolean enableNonAutoIdGenerateValue,boolean testPartition) {
        super(collectionName, ConnectType.MILVUS,policy);
        this.sinkPool = sinkPool;
        try {
            sinkMilvusClient = sinkPool.getConnection();
        } catch (InterruptedException e) {
            throw new IllegalStateException("目标milvus连接获取不到");
        }
        this.enableNonAutoIdGenerateValue = enableNonAutoIdGenerateValue;
        this.testPartition = testPartition;
    }

    @Override
    public boolean prepare(String source, MilvusCollectionMeta sourceMeta) {
        if(sourceMeta.getFieldTypes()==null || sourceMeta.getFieldTypes().isEmpty()){
            log.error("源{}字段列表为空", source);
            return false;
        }
        Optional<FieldTypeMeta> parimaryOptional = sourceMeta.getFieldTypes().stream()
                .filter(FieldTypeMeta::isPrimaryKey)
                .findFirst();
        if(parimaryOptional.isPresent()){
            isAutoID = parimaryOptional.get().isAutoID();
            primaryKey = parimaryOptional.get().getName();
            primaryDataType = parimaryOptional.get().getDataType();
        }else{
            log.error("没有找到源{}的主键字段",source);
            return false;
        }
        //获取向量字段
        sourceMeta.getFieldTypes().stream().filter(FieldTypeMeta::isVectorKey)
                .findFirst().ifPresent(fieldTypeMeta -> vectorKey = fieldTypeMeta.getName());

        boolean sinkExists = MilvusOperatorUtil.checkCollectionExist(getName(), sinkMilvusClient);
        boolean dropSinkCollectionSuccess = false;
        if (sinkExists) {
            if (PolicyType.COVER.equalsIgnoreCase(getPolicy())) {
                dropSinkCollectionSuccess = MilvusOperatorUtil.dropCollection(getName(), sinkMilvusClient);
                log.info(String.format("%s模式触发,删除已存在的目标collection【%s】完成", PolicyType.COVER, getName()));
            } else {
                DescCollResponseWrapper sinkDesc = MilvusOperatorUtil.getCollectionDescription(getName(), sinkMilvusClient);
                if (sinkDesc == null) {
                    log.error(String.format("获取不到目标collection【%s】的信息", getName()));
                    return false;
                }
                if (sinkDesc.getFields().size() != sourceMeta.getFieldTypes().size()) {
                    log.error("源【{}】和目标collection【{}】的字段列表不同，无法同步", source, getName());
                    return false;
                }
            }
        }
        if (!sinkExists || dropSinkCollectionSuccess) {
            //没有则创建
            CreateCollectionParam.Builder builder = CreateCollectionParam.newBuilder()
                    .withCollectionName(getName())
                    .withSchema(CollectionSchemaParam.newBuilder()
                            .withEnableDynamicField(sourceMeta.isEnableDynamicField())
                            .withFieldTypes(sourceMeta.toMilvusFieldTypes()).build())
                    .withDescription(sourceMeta.getCollectionDescription())
                    .withShardsNum(sourceMeta.getShardsNum());
            R<RpcStatus> collectionCreateStatus = sinkMilvusClient.createCollection(builder.build());
            if (R.success().getStatus().equals(collectionCreateStatus.getStatus())) {
                log.info("{}创建成功!", getName());
            } else {
                //skip 创建落地的collection失败
                log.error("{}创建失败:{}", getName(),collectionCreateStatus.getMessage());
                return false;
            }
        }

        return true;
    }

    @Override
    public long handler(String source,List<JSONObject> sourceJsonObjects) {
        if(sourceJsonObjects==null || sourceJsonObjects.isEmpty()) {
            return 0;
        }
        List<JsonObject> sourceGsonObjects = new ArrayList<>();
        //进行数据的预处理
        for (JSONObject jsonObject : sourceJsonObjects) {//处理主键
            handleAutoID(jsonObject, isAutoID, primaryKey, primaryDataType);
            //处理向量字段
            JSONArray vectorArray = jsonObject.getJSONArray(vectorKey);
            for(int i = 0; i < vectorArray.size(); i++){
                Object each = vectorArray.get(i);
                //需要将BigDecimal转为float，否则插入milvus时会出现报错Float vector field's value type must be List<Float>
                if(each instanceof BigDecimal){
                    vectorArray.set(i, ((BigDecimal) each).floatValue());
                }
            }

            JsonObject gsonObject = new JsonParser().parse(jsonObject.toJSONString()).getAsJsonObject();
            sourceGsonObjects.add(gsonObject);
        }

        InsertParam.Builder batchInsertParamBuilder  = InsertParam.newBuilder()
                .withCollectionName(getName()).withRows(sourceGsonObjects);
        
        String partition = null;
        if(testPartition){
            //TODO 这里犯懒了写死了逻辑，可以做成通用化拓展化的
            int randomIndex = random.nextInt(5);
            partition =  moniPartitionValueMap.get(randomIndex);
            if(!sinkMilvusClient.hasPartition(HasPartitionParam.newBuilder()
                    .withCollectionName(getName()).withPartitionName(partition).build()).getData()){
                if(sinkMilvusClient.createPartition(CreatePartitionParam.newBuilder()
                                .withCollectionName(getName()).withPartitionName(partition).build())
                        .getStatus() == R.Status.Success.getCode()){
                    log.info(String.format("目标collection【%s】创建分区%s完成",getName(),partition));
                }else{
                    log.error(String.format("目标collection【%s】创建分区%s失败!!",getName(),partition));
                }
            }
            batchInsertParamBuilder.withPartitionName(partition);
        }
        try{
            R<MutationResult> insertResponseR = sinkMilvusClient.insert(batchInsertParamBuilder.build());
            if(insertResponseR.getStatus() == R.Status.Success.getCode()){
                log.info(String.format("当前成功从源【%s】同步到目标collection【%s】%s条记录", source, getName(), sourceJsonObjects.size()));
                return sourceJsonObjects.size();
            }else{
                log.warn(String.format("当前从源【%s】同步到目标collection【%s】%s条记录出现问题:【%s】,准备逐条同步(同步速度将会下降)", source, getName(), sourceJsonObjects.size(),insertResponseR.getMessage()));
                int reInsertSuccessCount = 0;
                for(JsonObject sourceGsonObject:sourceGsonObjects){
                    InsertParam.Builder singleInsertParam =  InsertParam.newBuilder()
                            .withCollectionName(getName())
                            .withRows(Collections.singletonList(sourceGsonObject));
                    if(partition!=null){
                        singleInsertParam.withPartitionName(partition);
                    }
                    insertResponseR = sinkMilvusClient.insert(singleInsertParam.build());
                    if(insertResponseR.getStatus() == R.Status.Success.getCode()){
                        reInsertSuccessCount ++;
                    }else{
                        log.error(String.format("当前从源【%s】同步到目标collection【%s】1条记录出现问题:【%s】", source, getName(), insertResponseR.getMessage()));
                    }
                }
                log.info(String.format("当前成功从源【%s】逐条同步到目标collection【%s】共%s条记录", source, getName(), reInsertSuccessCount));
                return reInsertSuccessCount;
            }
        }catch (Exception e){
            log.error(String.format("当前从源【%s】同步到目标collection【%s】%s条记录失败", source, getName(), sourceJsonObjects.size()),e);
        }
        return 0;
    }

    @Override
    public void destroy() {
        if (sinkMilvusClient != null) {
            sinkPool.releaseConnection(sinkMilvusClient);
            sinkMilvusClient = null;
        }
    }

    private void handleAutoID(JSONObject jsonObject,boolean isAutoID,String primaryKey,int primaryDataType){
        if (isAutoID) {
            jsonObject.remove(primaryKey);
        } else {
            if (enableNonAutoIdGenerateValue) {
                if (primaryDataType == DataType.String.getNumber()
                        || primaryDataType == DataType.VarChar.getNumber()) {
                    jsonObject.put(primaryKey, IDGenerator.getNextIdStr());
                } else {
                    jsonObject.put(primaryKey, IDGenerator.getNextId());
                }
            }
        }
    }
}
